import cv2
import numpy as np

# 图像读取与预处理
img = cv2.imread('./images/lena.png')
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 定义Sobel卷积核
sobel_x = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float32)
sobel_y = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=np.float32)

# 执行卷积运算
gray_float = gray.astype(np.float32)
edge_x = cv2.filter2D(gray_float, -1, sobel_x, borderType=cv2.BORDER_CONSTANT)
edge_y = cv2.filter2D(gray_float, -1, sobel_y, borderType=cv2.BORDER_CONSTANT)

# 计算梯度幅值和方向
edge_mag = np.sqrt(edge_x**2 + edge_y**2)
edge_dir = np.arctan2(edge_y, edge_x)

# 归一化梯度图像用于显示
def normalize_image(img):
    normalized = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX)
    return normalized.astype(np.uint8)

edge_x_norm = normalize_image(edge_x)
edge_y_norm = normalize_image(edge_y)
edge_mag_norm = normalize_image(edge_mag)

# 绘制
# 确保所有图像都是3通道，便于拼接
if len(img.shape) == 3:
    img_display = img.copy()  # BGR格式，直接用于OpenCV显示
else:
    img_display = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

edge_x_display = cv2.cvtColor(edge_x_norm, cv2.COLOR_GRAY2BGR)
edge_y_display = cv2.cvtColor(edge_y_norm, cv2.COLOR_GRAY2BGR)
edge_mag_display = cv2.cvtColor(edge_mag_norm, cv2.COLOR_GRAY2BGR)


# 创建2x2的网格布局
h, w = img.shape[:2]
grid = np.zeros((h*2, w*2, 3), dtype=np.uint8)

# 将图像放入网格
grid[:h, :w] = img_display                # 左上：原始图像
grid[:h, w:2*w] = edge_mag_display        # 右上：边缘强度图
grid[h:2*h, :w] = edge_x_display          # 左下：水平梯度
grid[h:2*h, w:2*w] = edge_y_display       # 右下：垂直梯度

# 添加标题
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(grid, 'Original Image', (10, 30), font, 1, (255, 255, 255), 2)
cv2.putText(grid, 'Sobel Edge Detection', (w+10, 30), font, 1, (255, 255, 255), 2)
cv2.putText(grid, 'Sobel X Edge Detection', (10, h+30), font, 1, (255, 255, 255), 2)
cv2.putText(grid, 'Sobel X Edge Detection', (w+10, h+30), font, 1, (255, 255, 255), 2)

# 显示拼接后的图像
cv2.imshow('Sobel Edge Detection Results', grid)
cv2.waitKey(0)